package com.ksyun.campus.metaserver.services.impl;

import com.ksyun.campus.metaserver.config.MinfsProperties;
import com.ksyun.campus.metaserver.domain.*;
import com.ksyun.campus.metaserver.services.IMetaService;
import com.ksyun.campus.metaserver.storage.MetadataStore;
import com.ksyun.campus.metaserver.zk.ZooKeeperService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;

/**
 * 改进的MetaService实现，提供更好的三副本分配策略
 * 
 * 主要改进：
 * 1. 智能负载均衡：结合容量使用率和节点负载
 * 2. 机架感知：尽量将副本分散到不同机架
 * 3. 故障域分散：避免将副本放置在相关性高的节点上
 * 4. 动态权重调整：根据节点性能动态调整选择权重
 */
@Slf4j
@Service("improvedMetaService")
public class ImprovedMetaServiceImpl {
    
    @Autowired
    private MetadataStore metadataStore;
    
    @Autowired
    private ZooKeeperService zooKeeperService;
    
    @Autowired
    private MinfsProperties properties;
    
    // 负载均衡权重配置
    private static final double CAPACITY_WEIGHT = 0.4;      // 容量权重
    private static final double LOAD_WEIGHT = 0.3;          // 负载权重  
    private static final double RACK_WEIGHT = 0.2;          // 机架分散权重
    private static final double RANDOM_WEIGHT = 0.1;        // 随机因子权重
    
    /**
     * 改进的块分配方法
     * 使用多因素加权选择算法，实现更好的负载均衡和容错性
     */
    public List<BlockInfo> allocateBlocksImproved(String fileId, long bytes) throws Exception {
        // 获取文件inode
        Inode fileInode = metadataStore.getInode(fileId).orElseThrow();
        if (!fileInode.isFile()) {
            throw new IllegalArgumentException("Not a file: " + fileId);
        }
        
        // 获取集群配置
        ClusterConfig config = metadataStore.getClusterConfig().orElse(ClusterConfig.defaultConfig());
        long blockSize = config.getBlockSizeBytes();
        int numBlocks = (int) Math.ceil((double) bytes / blockSize);
        
        // 准备分配块
        List<BlockInfo> allocatedBlocks = new ArrayList<>();
        List<BlockInfo> existingBlocks = metadataStore.getFileBlocks(fileId);
        int startIdx = existingBlocks.size();
        
        // 获取所有健康的DataServer
        List<DataServerInfo> healthyDataServers = getHealthyDataServers();
        if (healthyDataServers.size() < properties.getCluster().getReplicationFactor()) {
            throw new RuntimeException("Insufficient healthy DataServers for replication factor " 
                    + properties.getCluster().getReplicationFactor());
        }
        
        // 为每个块分配副本
        for (int i = 0; i < numBlocks; i++) {
            int blkIdx = startIdx + i;
            long blkSize = Math.min(blockSize, bytes - (i * blockSize));
            
            // 使用改进的副本选择策略
            List<DataServerInfo> selectedDs = selectDataServersForBlockImproved(fileId, blkIdx, healthyDataServers);
            
            // 创建副本信息
            List<BlockInfo.ReplicaInfo> replicas = createReplicaInfos(selectedDs);
            
            // 创建块信息
            BlockInfo blockInfo = BlockInfo.builder()
                    .fileId(fileId)
                    .blkIdx(blkIdx)
                    .size(blkSize)
                    .replicas(replicas)
                    .build();
            
            // 存储块信息
            metadataStore.putBlock(fileId, blkIdx, blockInfo);
            allocatedBlocks.add(blockInfo);
            
            // 更新DataServer的预期负载（用于后续块的分配决策）
            updateExpectedLoad(selectedDs, blkSize);
        }
        
        log.info("Allocated {} blocks for file {} using improved strategy", numBlocks, fileId);
        return allocatedBlocks;
    }
    
    /**
     * 改进的DataServer选择策略
     * 综合考虑负载、容量、机架分布等多个因素
     */
    private List<DataServerInfo> selectDataServersForBlockImproved(String fileId, int blkIdx, 
                                                                  List<DataServerInfo> availableDs) {
        int replicationFactor = Math.min(properties.getCluster().getReplicationFactor(), availableDs.size());
        List<DataServerInfo> selectedDs = new ArrayList<>();
        Set<String> usedRacks = new HashSet<>();
        
        // 计算每个DataServer的选择权重
        Map<DataServerInfo, Double> weights = calculateDataServerWeights(availableDs, fileId, blkIdx);
        
        // 使用加权随机选择算法
        for (int i = 0; i < replicationFactor; i++) {
            DataServerInfo selected = selectDataServerByWeight(weights, usedRacks, selectedDs);
            if (selected != null) {
                selectedDs.add(selected);
                weights.remove(selected); // 移除已选择的节点
                
                // 记录已使用的机架
                String rack = getRackId(selected);
                if (rack != null) {
                    usedRacks.add(rack);
                }
            }
        }
        
        // 如果选择的节点数不足，使用备选策略
        if (selectedDs.size() < replicationFactor) {
            log.warn("Could not select {} replicas using improved strategy, falling back to simple selection", 
                    replicationFactor);
            return selectDataServersFallback(availableDs, replicationFactor, selectedDs);
        }
        
        log.debug("Selected DataServers for {}:{} - {}", fileId, blkIdx, 
                selectedDs.stream().map(DataServerInfo::getDsId).collect(Collectors.toList()));
        
        return selectedDs;
    }
    
    /**
     * 计算DataServer的选择权重
     * 综合考虑容量、负载、机架分布等因素
     */
    private Map<DataServerInfo, Double> calculateDataServerWeights(List<DataServerInfo> dataServers, 
                                                                  String fileId, int blkIdx) {
        Map<DataServerInfo, Double> weights = new HashMap<>();
        
        for (DataServerInfo ds : dataServers) {
            double weight = 0.0;
            
            // 1. 容量因子：剩余容量越多权重越高
            double capacityFactor = calculateCapacityFactor(ds);
            weight += CAPACITY_WEIGHT * capacityFactor;
            
            // 2. 负载因子：当前负载越低权重越高
            double loadFactor = calculateLoadFactor(ds);
            weight += LOAD_WEIGHT * loadFactor;
            
            // 3. 机架分散因子：不同机架的节点权重更高
            double rackFactor = calculateRackFactor(ds, fileId, blkIdx);
            weight += RACK_WEIGHT * rackFactor;
            
            // 4. 随机因子：增加一定的随机性
            double randomFactor = ThreadLocalRandom.current().nextDouble();
            weight += RANDOM_WEIGHT * randomFactor;
            
            weights.put(ds, weight);
        }
        
        return weights;
    }
    
    /**
     * 计算容量因子
     * 基于可用容量比例，容量越多权重越高
     */
    private double calculateCapacityFactor(DataServerInfo ds) {
        if (ds.getCapacityBytes() <= 0) {
            return 0.0;
        }
        
        double usageRatio = ds.getUsageRatio();
        // 使用反向的使用率，剩余容量越多权重越高
        return Math.max(0.0, 1.0 - usageRatio);
    }
    
    /**
     * 计算负载因子
     * 基于当前负载情况，负载越低权重越高
     */
    private double calculateLoadFactor(DataServerInfo ds) {
        // 这里可以根据实际的负载指标进行计算
        // 暂时使用容量使用率作为负载指标
        double usageRatio = ds.getUsageRatio();
        
        // 使用平滑的权重函数，避免极端值
        if (usageRatio < 0.5) {
            return 1.0; // 低负载，高权重
        } else if (usageRatio < 0.8) {
            return 0.7; // 中等负载，中等权重
        } else {
            return 0.3; // 高负载，低权重
        }
    }
    
    /**
     * 计算机架分散因子
     * 基于一致性哈希和机架信息，鼓励跨机架分布
     */
    private double calculateRackFactor(DataServerInfo ds, String fileId, int blkIdx) {
        // 使用一致性哈希确保相同文件块的副本分布相对稳定
        long hash = Math.abs(Objects.hash(fileId, blkIdx, ds.getDsId()));
        double hashFactor = (hash % 1000) / 1000.0;
        
        // 机架信息可以从labels中获取
        String rack = getRackId(ds);
        if (rack != null) {
            // 如果有机架信息，增加跨机架分布的权重
            long rackHash = Math.abs(Objects.hash(fileId, blkIdx, rack));
            double rackHashFactor = (rackHash % 1000) / 1000.0;
            return (hashFactor + rackHashFactor) / 2.0;
        }
        
        return hashFactor;
    }
    
    /**
     * 根据权重选择DataServer
     * 使用加权随机选择，同时考虑机架分散性
     */
    private DataServerInfo selectDataServerByWeight(Map<DataServerInfo, Double> weights, 
                                                   Set<String> usedRacks, 
                                                   List<DataServerInfo> alreadySelected) {
        if (weights.isEmpty()) {
            return null;
        }
        
        // 首先尝试选择不同机架的节点
        Map<DataServerInfo, Double> preferredNodes = weights.entrySet().stream()
                .filter(entry -> {
                    String rack = getRackId(entry.getKey());
                    return rack == null || !usedRacks.contains(rack);
                })
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        
        // 如果没有不同机架的节点，使用所有可用节点
        Map<DataServerInfo, Double> candidateNodes = preferredNodes.isEmpty() ? weights : preferredNodes;
        
        // 计算总权重
        double totalWeight = candidateNodes.values().stream().mapToDouble(Double::doubleValue).sum();
        if (totalWeight <= 0) {
            // 如果所有权重都为0，随机选择一个
            List<DataServerInfo> candidates = new ArrayList<>(candidateNodes.keySet());
            return candidates.get(ThreadLocalRandom.current().nextInt(candidates.size()));
        }
        
        // 加权随机选择
        double randomValue = ThreadLocalRandom.current().nextDouble() * totalWeight;
        double currentWeight = 0.0;
        
        for (Map.Entry<DataServerInfo, Double> entry : candidateNodes.entrySet()) {
            currentWeight += entry.getValue();
            if (currentWeight >= randomValue) {
                return entry.getKey();
            }
        }
        
        // 备选方案：返回权重最高的节点
        return candidateNodes.entrySet().stream()
                .max(Map.Entry.comparingByValue())
                .map(Map.Entry::getKey)
                .orElse(null);
    }
    
    /**
     * 获取节点的机架ID
     */
    private String getRackId(DataServerInfo ds) {
        if (ds.getLabels() != null) {
            return ds.getLabels().get("rack");
        }
        return null;
    }
    
    /**
     * 备选的DataServer选择策略
     * 当主要策略无法选择足够节点时使用
     */
    private List<DataServerInfo> selectDataServersFallback(List<DataServerInfo> availableDs, 
                                                          int replicationFactor,
                                                          List<DataServerInfo> alreadySelected) {
        List<DataServerInfo> result = new ArrayList<>(alreadySelected);
        Set<String> selectedIds = alreadySelected.stream()
                .map(DataServerInfo::getDsId)
                .collect(Collectors.toSet());
        
        // 按使用率排序，优先选择使用率较低的节点
        List<DataServerInfo> sortedDs = availableDs.stream()
                .filter(ds -> !selectedIds.contains(ds.getDsId()))
                .sorted(Comparator.comparingDouble(DataServerInfo::getUsageRatio))
                .collect(Collectors.toList());
        
        // 补充选择节点直到满足副本数要求
        for (DataServerInfo ds : sortedDs) {
            if (result.size() >= replicationFactor) {
                break;
            }
            result.add(ds);
        }
        
        return result;
    }
    
    /**
     * 创建副本信息列表
     */
    private List<BlockInfo.ReplicaInfo> createReplicaInfos(List<DataServerInfo> selectedDs) {
        List<BlockInfo.ReplicaInfo> replicas = new ArrayList<>();
        
        for (int i = 0; i < selectedDs.size(); i++) {
            DataServerInfo ds = selectedDs.get(i);
            BlockInfo.ReplicaInfo.ReplicaRole role = (i == 0) ? 
                    BlockInfo.ReplicaInfo.ReplicaRole.PRIMARY : 
                    BlockInfo.ReplicaInfo.ReplicaRole.REPLICA;
            
            replicas.add(BlockInfo.ReplicaInfo.builder()
                    .dsId(ds.getDsId())
                    .host(ds.getHost())
                    .port(ds.getPort())
                    .role(role)
                    .build());
        }
        
        return replicas;
    }
    
    /**
     * 更新预期负载
     * 用于在同一次分配中的后续决策
     */
    private void updateExpectedLoad(List<DataServerInfo> selectedDs, long blockSize) {
        for (DataServerInfo ds : selectedDs) {
            // 临时增加使用量，用于后续块分配的负载计算
            ds.setUsedBytes(ds.getUsedBytes() + blockSize);
        }
    }
    
    /**
     * 获取健康的DataServer列表
     */
    private List<DataServerInfo> getHealthyDataServers() {
        return zooKeeperService.getCurrentDataServers().stream()
                .filter(ds -> ds.isHealthy(properties.getHeartbeat().getTimeoutSeconds()))
                .filter(ds -> ds.getStatus() == DataServerInfo.DataServerStatus.ACTIVE)
                .collect(Collectors.toList());
    }
    
    /**
     * 验证副本分配是否满足要求
     */
    public boolean validateReplicaPlacement(List<BlockInfo.ReplicaInfo> replicas) {
        if (replicas.size() < properties.getCluster().getReplicationFactor()) {
            log.warn("Replica count {} is less than required replication factor {}", 
                    replicas.size(), properties.getCluster().getReplicationFactor());
            return false;
        }
        
        // 检查是否有重复的DataServer
        Set<String> dsIds = replicas.stream()
                .map(BlockInfo.ReplicaInfo::getDsId)
                .collect(Collectors.toSet());
        
        if (dsIds.size() != replicas.size()) {
            log.warn("Duplicate DataServer found in replica placement");
            return false;
        }
        
        // 检查是否有PRIMARY副本
        boolean hasPrimary = replicas.stream()
                .anyMatch(r -> r.getRole() == BlockInfo.ReplicaInfo.ReplicaRole.PRIMARY);
        
        if (!hasPrimary) {
            log.warn("No PRIMARY replica found in replica placement");
            return false;
        }
        
        return true;
    }
}
